//***************************************************************************************
// LightingUtil.hlsl by Frank Luna (C) 2015 All Rights Reserved.
//
// Contains API for shader lighting.
//***************************************************************************************
#ifndef __SHADER_LIGHTUTIL__
#define __SHADER_LIGHTUTIL__
#include "Utlis.hlsl"
#define MaxLights 16

struct Light
{
    float3 Strength;
    float FalloffStart; // point/spot light only
    float3 Direction;   // directional/spot light only
    float FalloffEnd;   // point/spot light only
    float3 Position;    // point light only
    float SpotPower;    // spot light only
    
    // for PBR lighting
	float3 Color; // All light
	float Intensity; // All light
	float3 _Position; // Point/Spot light only
	float Range; // Point/Spot light only
	float3 _Direction; // Directional/Spot light only
	float SpotRadius; // Spot light only
	float2 SpotAngles; // Spot light only
	uint LightType;
	int ShadowMapIdx; //Shadowmap Per Light
    float4x4 LightProj;//projection
	float4x4 LightView;//view
    
};

struct Material
{
    float4 DiffuseAlbedo;
    float3 FresnelR0;
    float Shininess;
};

float CalcAttenuation(float d, float falloffStart, float falloffEnd)
{
    // Linear falloff.
    return saturate((falloffEnd-d) / (falloffEnd - falloffStart));
}

// Schlick gives an approximation to Fresnel reflectance (see pg. 233 "Real-Time Rendering 3rd Ed.").
// R0 = ( (n-1)/(n+1) )^2, where n is the index of refraction.
float3 SchlickFresnel(float3 R0, float3 normal, float3 lightVec)
{
    float cosIncidentAngle = saturate(dot(normal, lightVec));

    float f0 = 1.0f - cosIncidentAngle;
    float3 reflectPercent = R0 + (1.0f - R0)*(f0*f0*f0*f0*f0);

    return reflectPercent;
}

float3 BlinnPhong(float3 lightStrength, float3 lightVec, float3 normal, float3 toEye, Material mat)
{
    const float m = mat.Shininess * 256.0f;
    float3 halfVec = normalize(toEye + lightVec);

    float roughnessFactor = (m + 8.0f)*pow(max(dot(halfVec, normal), 0.0f), m) / 8.0f;
    float3 fresnelFactor = SchlickFresnel(mat.FresnelR0, halfVec, lightVec);

    float3 specAlbedo = fresnelFactor*roughnessFactor;

    // Our spec formula goes outside [0,1] range, but we are 
    // doing LDR rendering.  So scale it down a bit.
    specAlbedo = specAlbedo / (specAlbedo + 1.0f);

    return (mat.DiffuseAlbedo.rgb + specAlbedo) * lightStrength;
}

//---------------------------------------------------------------------------------------
// Evaluates the lighting equation for directional lights.
//---------------------------------------------------------------------------------------
float3 ComputeDirectionalLight(Light L, Material mat, float3 normal, float3 toEye)
{
    // The light vector aims opposite the direction the light rays travel.
    float3 lightVec = normalize(-L.Direction);
    
    // Scale light down by Lambert's cosine law.
    float ndotl = max(dot(lightVec, normal), 0.0f);
    float3 lightStrength = L.Strength * ndotl;

    return BlinnPhong(lightStrength, lightVec, normal, toEye, mat);
}

//---------------------------------------------------------------------------------------
// Evaluates the lighting equation for point lights.
//---------------------------------------------------------------------------------------
float3 ComputePointLight(Light L, Material mat, float3 pos, float3 normal, float3 toEye)
{
    // The vector from the surface to the light.
    float3 lightVec = L.Position - pos;

    // The distance from surface to light.
    float d = length(lightVec);

    // Range test.
    if(d > L.FalloffEnd)
        return float3(0,0,0);

    // Normalize the light vector.
    lightVec /= d;

    // Scale light down by Lambert's cosine law.
    float ndotl = max(dot(lightVec, normal), 0.0f);
    float3 lightStrength = L.Strength * ndotl;

    // Attenuate light by distance.
    float att = CalcAttenuation(d, L.FalloffStart, L.FalloffEnd);
    lightStrength *= att;

    return BlinnPhong(lightStrength, lightVec, normal, toEye, mat);
}

//---------------------------------------------------------------------------------------
// Evaluates the lighting equation for spot lights.
//---------------------------------------------------------------------------------------
float3 ComputeSpotLight(Light L, Material mat, float3 pos, float3 normal, float3 toEye)
{
    // The vector from the surface to the light.
    float3 lightVec = L.Position - pos;

    // The distance from surface to light.
    float d = length(lightVec);

    // Range test.
    if(d > L.FalloffEnd)
        return float3(0, 0, 0);

    // Normalize the light vector.
    lightVec /= d;

    // Scale light down by Lambert's cosine law.
    float ndotl = max(dot(lightVec, normal), 0.0f);
    float3 lightStrength = L.Strength * ndotl;

    // Attenuate light by distance.
    float att = CalcAttenuation(d, L.FalloffStart, L.FalloffEnd);
    lightStrength *= att;

    // Scale by spotlight
    float spotFactor = pow(max(dot(-lightVec, L.Direction), 0.0f), L.SpotPower);
    lightStrength *= spotFactor;

    return BlinnPhong(lightStrength, lightVec, normal, toEye, mat);
}

float4 ComputeLighting(StructuredBuffer<Light> gLights, Material mat,
                       float3 pos, float3 normal, float3 toEye,
                       float3 shadowFactor, int gLightCount)
{
    float3 result = 0.0f;

    int i = 0;

    [unroll(10)]
    for (int j = 0; j < gLightCount; ++j)
    {
        if (gLights[j].LightType == 0)
        {
            result += shadowFactor[i] * ComputeDirectionalLight(gLights[j], mat, normal, toEye) * gLights[j].Color;
            i++;
        }
        else if (gLights[j].LightType == 1)
        {
            result += ComputeSpotLight(gLights[j], mat, pos, normal, toEye) * gLights[j].Color;
        }else if (gLights[j].LightType == 2)
        {
            result += ComputePointLight(gLights[j], mat, pos, normal, toEye) * gLights[j].Color;
        }
    }
    return float4(result, 0.0f);
}

//----------------------------------------------------------------------------------------------
// PBR
//----------------------------------------------------------------------------------------------
static const float F0_DIELECTRIC = 0.04f;

// Lambert diffuse
float3 LambertDiffuse(float3 DiffuseColor)
{
	return DiffuseColor * (1 / PI);
}

// GGX (Trowbridge-Reitz)
float GGX(float a2, float NoH)
{
	float NoH2 = NoH * NoH;
	float d = NoH2 * (a2 - 1.0f) + 1.0f;

	return a2 / (PI * d * d);
}

// Fresnel, Schlick approx.
float3 FresnelSchlick(float3 F0, float VoH)
{
	return F0 + (1 - F0) * exp2((-5.55473 * VoH - 6.98316) * VoH);
}

// Schlick-GGX
float GeometrySchlickGGX(float Roughness, float NoV)
{
	float k = pow(Roughness + 1, 2) / 8.0f;

	return NoV / (NoV * (1 - k) + k);
}

// Cook-Torrance Specular
float3 SpecularGGX(float3 N, float3 L, float3 V, float Roughness, float3 F0)
{
	float3 H = normalize(V + L);
    
	float NoL = saturate(dot(N, L));
	float NoV = saturate(dot(N, V));
	float VoH = saturate(dot(V, H));
	float NoH = saturate(dot(N, H));

	float a2 = Pow4(Roughness);
	float D = GGX(a2, NoH);
	float3 F = FresnelSchlick(F0, VoH);
	float G = GeometrySchlickGGX(Roughness, NoV) * GeometrySchlickGGX(Roughness, NoL);

	return (D * G * F) / (4 * max(NoL * NoV, 0.01f)); // 0.01 is added to prevent division by 0
}

float3 DefaultBRDF(float3 LightDir, float3 Normal, float3 ViewDir, float Roughness, float Metallic, float3 BaseColor)
{
	float3 F0 = lerp(F0_DIELECTRIC.rrr, BaseColor.rgb, Metallic);
    
    // Base color remapping
	float3 DiffuseColor = (1.0 - Metallic) * BaseColor; // Metallic surfaces have no diffuse reflections
    
	float3 DiffuseBRDF = LambertDiffuse(DiffuseColor);
	float3 SpecularBRDF = SpecularGGX(Normal, LightDir, ViewDir, Roughness, F0);
    
	return DiffuseBRDF + SpecularBRDF;
}

float3 DirectLighting(float3 Radiance, float3 LightDir, float3 Normal, float3 ViewDir, float Roughness, float Metallic, float3 BaseColor, float ShadowFactor)
{
    float3 BRDF = DefaultBRDF(LightDir, Normal, ViewDir, Roughness, Metallic, BaseColor);
    
    float NoL = saturate(dot(Normal, LightDir));
    
    return Radiance * BRDF * NoL * ShadowFactor;
}

float CalcDistanceAttenuation(float3 WorldPos, float3 LightPos, float Range)
{
    float3 WorldLightVector = LightPos - WorldPos;
    float DistanceSqr = dot(WorldLightVector, WorldLightVector);
    
    float DistanceAttenuation = 1 / (DistanceSqr + 1);
    
    float InvRange = rcp(max(Range, 0.001f));
    float LightRadiusMask = Square(saturate(1 - Square(DistanceSqr * InvRange * InvRange)));
    DistanceAttenuation *= LightRadiusMask;
    
    return DistanceAttenuation;
}

/**
 * LightDir is the direction from shader point to spot light.
 * SpotDirection is the direction of the spot light.
 * SpotAngles.x is CosOuterCone, SpotAngles.y is InvCosConeDifference. 
*/
float SpotAttenuation(float3 LightDir, float3 SpotDirection, float2 SpotAngles)
{
    float ConeAngleFalloff = Square(saturate((dot(-LightDir, SpotDirection) - SpotAngles.x) * SpotAngles.y));
    return ConeAngleFalloff;
}
#endif //__SHADER_LIGHTUTIL__

